package com.pai4j.aigc.controller;

import com.pai4j.aigc.llm.entity.LLMUsageEntity;
import com.pai4j.aigc.llm.repository.dao.ILLMUsageDAO;
import com.pai4j.common.bean.PAIResponseBean;
import com.pai4j.common.bean.PAIResponseBeanUtil;
import com.pai4j.log.annotation.OperationLog;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.http.ResponseEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

import jakarta.persistence.criteria.Predicate;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.HashMap;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/aigc/admin/llm/usage")
public class LLMUsageController {

    @Autowired
    private ILLMUsageDAO usageDAO;

    @GetMapping("/list")
    @OperationLog(module = "AIGC管理", operation = "LLM用量查询", description = "按租户分页查询LLM用量记录（支持多条件过滤）", recordParams = false)
    public PAIResponseBean<Page<LLMUsageEntity>> list(@RequestParam String tenantId,
                                                      @RequestParam(required = false) String userId,
                                                      @RequestParam(required = false) Long modelId,
                                                      @RequestParam(required = false) Boolean success,
                                                      @RequestParam(required = false) String startTime,
                                                      @RequestParam(required = false) String endTime,
                                                      @RequestParam(defaultValue = "0") int page,
                                                      @RequestParam(defaultValue = "20") int size) {
        Pageable pageable = PageRequest.of(page, size);
        Specification<LLMUsageEntity> spec = (root, query, cb) -> {
            List<Predicate> predicates = new ArrayList<>();
            predicates.add(cb.equal(root.get("tenantId"), tenantId));
            if (userId != null && !userId.isEmpty()) {
                predicates.add(cb.equal(root.get("userId"), userId));
            }
            if (modelId != null && modelId > 0) {
                predicates.add(cb.equal(root.get("modelId"), modelId));
            }
            if (success != null) {
                predicates.add(cb.equal(root.get("success"), success ? 1 : 0));
            }
            if (startTime != null && !startTime.isEmpty()) {
                predicates.add(cb.greaterThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(startTime)));
            }
            if (endTime != null && !endTime.isEmpty()) {
                predicates.add(cb.lessThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(endTime)));
            }
            return cb.and(predicates.toArray(new Predicate[0]));
        };
        Page<LLMUsageEntity> result = usageDAO.findAll(spec, pageable);
        return PAIResponseBeanUtil.success(result);
    }

    @GetMapping("/stats/group")
    @OperationLog(module = "AIGC管理", operation = "LLM用量统计分组", description = "按 model 或 user 维度聚合统计", recordParams = false)
    public PAIResponseBean<List<Map<String, Object>>> groupedStats(@RequestParam String tenantId,
                                                                   @RequestParam String groupBy, // model | user
                                                                   @RequestParam(required = false) String userId,
                                                                   @RequestParam(required = false) Long modelId,
                                                                   @RequestParam(required = false) String startTime,
                                                                   @RequestParam(required = false) String endTime) {
        Specification<LLMUsageEntity> spec = (root, query, cb) -> {
            List<Predicate> predicates = new ArrayList<>();
            predicates.add(cb.equal(root.get("tenantId"), tenantId));
            if (userId != null && !userId.isEmpty()) {
                predicates.add(cb.equal(root.get("userId"), userId));
            }
            if (modelId != null && modelId > 0) {
                predicates.add(cb.equal(root.get("modelId"), modelId));
            }
            if (startTime != null && !startTime.isEmpty()) {
                predicates.add(cb.greaterThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(startTime)));
            }
            if (endTime != null && !endTime.isEmpty()) {
                predicates.add(cb.lessThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(endTime)));
            }
            return cb.and(predicates.toArray(new Predicate[0]));
        };

        List<LLMUsageEntity> list = usageDAO.findAll(spec);
        Map<String, List<LLMUsageEntity>> grouped = new HashMap<>();
        if ("model".equalsIgnoreCase(groupBy)) {
            grouped = list.stream().collect(Collectors.groupingBy(e -> String.valueOf(e.getModelId())));
        } else if ("user".equalsIgnoreCase(groupBy)) {
            grouped = list.stream().collect(Collectors.groupingBy(LLMUsageEntity::getUserId));
        } else {
            throw new IllegalArgumentException("groupBy must be 'model' or 'user'");
        }

        List<Map<String, Object>> rows = new ArrayList<>();
        for (Map.Entry<String, List<LLMUsageEntity>> entry : grouped.entrySet()) {
            List<LLMUsageEntity> g = entry.getValue();
            long total = g.size();
            long totalSuccess = g.stream().filter(e -> e.getSuccess() != null && e.getSuccess() == 1).count();
            int sumPrompt = g.stream().map(e -> e.getTokensPrompt() == null ? 0 : e.getTokensPrompt()).reduce(0, Integer::sum);
            int sumCompletion = g.stream().map(e -> e.getTokensCompletion() == null ? 0 : e.getTokensCompletion()).reduce(0, Integer::sum);
            int sumTotal = g.stream().map(e -> e.getTokensTotal() == null ? 0 : e.getTokensTotal()).reduce(0, Integer::sum);
            int sumCostCents = g.stream().map(e -> e.getCostCents() == null ? 0 : e.getCostCents()).reduce(0, Integer::sum);
            double avgLatency = g.isEmpty() ? 0 : g.stream().map(e -> e.getLatencyMs() == null ? 0 : e.getLatencyMs()).collect(Collectors.averagingInt(i -> i));

            Map<String, Object> row = new HashMap<>();
            row.put("group_key", entry.getKey());
            row.put("total", total);
            row.put("success_total", totalSuccess);
            row.put("prompt_tokens", sumPrompt);
            row.put("completion_tokens", sumCompletion);
            row.put("total_tokens", sumTotal);
            row.put("cost_cents", sumCostCents);
            row.put("avg_latency_ms", avgLatency);
            rows.add(row);
        }
        return PAIResponseBeanUtil.success(rows);
    }

    @GetMapping("/export")
    @OperationLog(module = "AIGC管理", operation = "LLM用量导出", description = "导出用量明细或分组统计为CSV", recordParams = false)
    public ResponseEntity<byte[]> export(@RequestParam String tenantId,
                                         @RequestParam(defaultValue = "csv") String format, // csv|excel
                                         @RequestParam(defaultValue = "none") String groupBy, // none|model|user
                                         @RequestParam(required = false) String userId,
                                         @RequestParam(required = false) Long modelId,
                                         @RequestParam(required = false) String startTime,
                                         @RequestParam(required = false) String endTime) {
        // 查询数据
        Specification<LLMUsageEntity> spec = (root, query, cb) -> {
            List<Predicate> predicates = new ArrayList<>();
            predicates.add(cb.equal(root.get("tenantId"), tenantId));
            if (userId != null && !userId.isEmpty()) {
                predicates.add(cb.equal(root.get("userId"), userId));
            }
            if (modelId != null && modelId > 0) {
                predicates.add(cb.equal(root.get("modelId"), modelId));
            }
            if (startTime != null && !startTime.isEmpty()) {
                predicates.add(cb.greaterThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(startTime)));
            }
            if (endTime != null && !endTime.isEmpty()) {
                predicates.add(cb.lessThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(endTime)));
            }
            return cb.and(predicates.toArray(new Predicate[0]));
        };

        String lower = format.toLowerCase();
        if ("excel".equals(lower)) {
            // 生成简单HTML表格，Excel可直接打开
            StringBuilder html = new StringBuilder();
            html.append("<html><head><meta charset='UTF-8'></head><body><table border='1'>");
            if ("none".equalsIgnoreCase(groupBy)) {
                List<LLMUsageEntity> list = usageDAO.findAll(spec);
                html.append("<tr><th>request_id</th><th>tenant_id</th><th>user_id</th><th>model_id</th><th>provider</th><th>model_code</th><th>success</th><th>prompt_tokens</th><th>completion_tokens</th><th>total_tokens</th><th>cost_cents</th><th>latency_ms</th><th>create_time</th></tr>");
                for (LLMUsageEntity e : list) {
                    html.append("<tr>")
                        .append(td(e.getRequestId()))
                        .append(td(e.getTenantId()))
                        .append(td(e.getUserId()))
                        .append(td(String.valueOf(e.getModelId())))
                        .append(td(e.getProvider()))
                        .append(td(e.getModelCode()))
                        .append(td(String.valueOf(e.getSuccess())))
                        .append(td(n(e.getTokensPrompt())))
                        .append(td(n(e.getTokensCompletion())))
                        .append(td(n(e.getTokensTotal())))
                        .append(td(n(e.getCostCents())))
                        .append(td(n(e.getLatencyMs())))
                        .append(td(String.valueOf(e.getCreateDate())))
                        .append("</tr>");
                }
            } else {
                List<Map<String, Object>> rows = groupedStats(tenantId, groupBy, userId, modelId, startTime, endTime).getData();
                html.append("<tr><th>group_key</th><th>total</th><th>success_total</th><th>prompt_tokens</th><th>completion_tokens</th><th>total_tokens</th><th>cost_cents</th><th>avg_latency_ms</th></tr>");
                for (Map<String, Object> r : rows) {
                    html.append("<tr>")
                        .append(td(String.valueOf(r.get("group_key"))))
                        .append(td(String.valueOf(r.get("total"))))
                        .append(td(String.valueOf(r.get("success_total"))))
                        .append(td(String.valueOf(r.get("prompt_tokens"))))
                        .append(td(String.valueOf(r.get("completion_tokens"))))
                        .append(td(String.valueOf(r.get("total_tokens"))))
                        .append(td(String.valueOf(r.get("cost_cents"))))
                        .append(td(String.valueOf(r.get("avg_latency_ms"))))
                        .append("</tr>");
                }
            }
            html.append("</table></body></html>");
            byte[] bytes = html.toString().getBytes(java.nio.charset.StandardCharsets.UTF_8);
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.parseMediaType("application/vnd.ms-excel"));
            headers.set(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=llm_usage_" + groupBy + ".xls");
            return ResponseEntity.ok().headers(headers).body(bytes);
        } else {
            // CSV 输出
            StringBuilder csv = new StringBuilder();
            if ("none".equalsIgnoreCase(groupBy)) {
                // 明细导出
                List<LLMUsageEntity> list = usageDAO.findAll(spec);
                csv.append("request_id,tenant_id,user_id,model_id,provider,model_code,success,prompt_tokens,completion_tokens,total_tokens,cost_cents,latency_ms,create_time\n");
                for (LLMUsageEntity e : list) {
                    csv.append(s(e.getRequestId())).append(',')
                       .append(s(e.getTenantId())).append(',')
                       .append(s(e.getUserId())).append(',')
                       .append(e.getModelId()).append(',')
                       .append(s(e.getProvider())).append(',')
                       .append(s(e.getModelCode())).append(',')
                       .append(e.getSuccess()).append(',')
                       .append(n(e.getTokensPrompt())).append(',')
                       .append(n(e.getTokensCompletion())).append(',')
                       .append(n(e.getTokensTotal())).append(',')
                       .append(n(e.getCostCents())).append(',')
                       .append(n(e.getLatencyMs())).append(',')
                       .append(s(String.valueOf(e.getCreateDate()))).append('\n');
                }
            } else {
                // 分组导出（model|user）
                List<Map<String, Object>> rows = groupedStats(tenantId, groupBy, userId, modelId, startTime, endTime).getData();
                csv.append("group_key,total,success_total,prompt_tokens,completion_tokens,total_tokens,cost_cents,avg_latency_ms\n");
                for (Map<String, Object> r : rows) {
                    csv.append(s(String.valueOf(r.get("group_key")))).append(',')
                       .append(s(String.valueOf(r.get("total")))).append(',')
                       .append(s(String.valueOf(r.get("success_total")))).append(',')
                       .append(s(String.valueOf(r.get("prompt_tokens")))).append(',')
                       .append(s(String.valueOf(r.get("completion_tokens")))).append(',')
                       .append(s(String.valueOf(r.get("total_tokens")))).append(',')
                       .append(s(String.valueOf(r.get("cost_cents")))).append(',')
                       .append(s(String.valueOf(r.get("avg_latency_ms")))).append('\n');
                }
            }
            byte[] bytes = csv.toString().getBytes(java.nio.charset.StandardCharsets.UTF_8);
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.TEXT_PLAIN);
            headers.set(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=llm_usage_" + groupBy + ".csv");
            return ResponseEntity.ok().headers(headers).body(bytes);
        }
    }

    private String s(String v) {
        if (v == null) return "";
        // 简单CSV转义
        String val = v.replace("\"", "\"\"");
        if (val.contains(",") || val.contains("\n") || val.contains("\r")) {
            return "\"" + val + "\"";
        }
        return val;
    }

    private String s(Object v) { return v == null ? "" : s(String.valueOf(v)); }

    private String n(Integer v) { return v == null ? "" : String.valueOf(v); }

    private String td(String v) {
        if (v == null) v = "";
        String esc = v.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;");
        return "<td>" + esc + "</td>";
    }

    @GetMapping("/stats")
    @OperationLog(module = "AIGC管理", operation = "LLM用量统计", description = "汇总某段时间内的用量与成本", recordParams = false)
    public PAIResponseBean<Map<String, Object>> stats(@RequestParam String tenantId,
                                                      @RequestParam(required = false) String userId,
                                                      @RequestParam(required = false) Long modelId,
                                                      @RequestParam(required = false) String startTime,
                                                      @RequestParam(required = false) String endTime) {
        Specification<LLMUsageEntity> spec = (root, query, cb) -> {
            List<Predicate> predicates = new ArrayList<>();
            predicates.add(cb.equal(root.get("tenantId"), tenantId));
            if (userId != null && !userId.isEmpty()) {
                predicates.add(cb.equal(root.get("userId"), userId));
            }
            if (modelId != null && modelId > 0) {
                predicates.add(cb.equal(root.get("modelId"), modelId));
            }
            if (startTime != null && !startTime.isEmpty()) {
                predicates.add(cb.greaterThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(startTime)));
            }
            if (endTime != null && !endTime.isEmpty()) {
                predicates.add(cb.lessThanOrEqualTo(root.get("createDate"), LocalDateTime.parse(endTime)));
            }
            return cb.and(predicates.toArray(new Predicate[0]));
        };

        List<LLMUsageEntity> list = usageDAO.findAll(spec);
        long total = list.size();
        long totalSuccess = list.stream().filter(e -> e.getSuccess() != null && e.getSuccess() == 1).count();
        int sumPrompt = list.stream().map(e -> e.getTokensPrompt() == null ? 0 : e.getTokensPrompt()).reduce(0, Integer::sum);
        int sumCompletion = list.stream().map(e -> e.getTokensCompletion() == null ? 0 : e.getTokensCompletion()).reduce(0, Integer::sum);
        int sumTotal = list.stream().map(e -> e.getTokensTotal() == null ? 0 : e.getTokensTotal()).reduce(0, Integer::sum);
        int sumCostCents = list.stream().map(e -> e.getCostCents() == null ? 0 : e.getCostCents()).reduce(0, Integer::sum);
        double avgLatency = list.isEmpty() ? 0 : list.stream().map(e -> e.getLatencyMs() == null ? 0 : e.getLatencyMs()).collect(Collectors.averagingInt(i -> i));

        Map<String, Object> result = new HashMap<>();
        result.put("total", total);
        result.put("success_total", totalSuccess);
        result.put("prompt_tokens", sumPrompt);
        result.put("completion_tokens", sumCompletion);
        result.put("total_tokens", sumTotal);
        result.put("cost_cents", sumCostCents);
        result.put("avg_latency_ms", avgLatency);
        return PAIResponseBeanUtil.success(result);
    }
}
